{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import math\n",
    "import numpy as np\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "from keras_bert import load_vocabulary, load_trained_model_from_checkpoint, Tokenizer, get_checkpoint_paths\n",
    "\n",
    "import keras.backend as K\n",
    "from keras.layers import Input, Dense, Lambda, Multiply, Masking, Concatenate\n",
    "from keras.models import Model\n",
    "from keras.preprocessing.sequence import pad_sequences\n",
    "from keras.callbacks import Callback, ModelCheckpoint\n",
    "from keras.utils.data_utils import Sequence\n",
    "from keras.utils import multi_gpu_model\n",
    "\n",
    "from nl2sql.utils import read_data, read_tables, SQL, MultiSentenceTokenizer, Query, Question, Table\n",
    "from nl2sql.utils.optimizer import RAdam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_table_file = '../data/train.tables.json'\n",
    "train_data_file = '../data/train.json'\n",
    "\n",
    "val_table_file = '../data/val.tables.json'\n",
    "val_data_file = '../data/val.json'\n",
    "\n",
    "test_table_file = '../data/test.tables.json'\n",
    "test_data_file = '../data/test.json'\n",
    "\n",
    "# Download pretrained BERT model from https://github.com/ymcui/Chinese-BERT-wwm\n",
    "bert_model_path = '../model/chinese_wwm_L-12_H-768_A-12'\n",
    "paths = get_checkpoint_paths(bert_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tables = read_tables(train_table_file)\n",
    "train_data = read_data(train_data_file, train_tables)\n",
    "\n",
    "val_tables = read_tables(val_table_file)\n",
    "val_data = read_data(val_data_file, val_tables)\n",
    "\n",
    "test_tables = read_tables(test_table_file)\n",
    "test_data = read_data(test_data_file, test_tables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_query = train_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "二零一九年第四周大黄蜂和密室逃生这两部影片的票房总占比是多少呀"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_query.question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "sel: [2]<br>agg: ['SUM']<br>cond_conn_op: 'or'<br>conds: [[0, '==', '大黄蜂'], [0, '==', '密室逃生']]"
      ],
      "text/plain": [
       "sel: [2]\n",
       "agg: ['SUM']\n",
       "cond_conn_op: 'or'\n",
       "conds: [[0, '==', '大黄蜂'], [0, '==', '密室逃生']]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_query.sql"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>影片名称</th>\n",
       "      <th>周票房（万）</th>\n",
       "      <th>票房占比（%）</th>\n",
       "      <th>场均人次</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>死侍2：我爱我家</td>\n",
       "      <td>10637.3</td>\n",
       "      <td>25.8</td>\n",
       "      <td>5.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>白蛇：缘起</td>\n",
       "      <td>10503.8</td>\n",
       "      <td>25.4</td>\n",
       "      <td>7.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>大黄蜂</td>\n",
       "      <td>6426.6</td>\n",
       "      <td>15.6</td>\n",
       "      <td>6.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>密室逃生</td>\n",
       "      <td>5841.4</td>\n",
       "      <td>14.2</td>\n",
       "      <td>6.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>“大”人物</td>\n",
       "      <td>3322.9</td>\n",
       "      <td>8.1</td>\n",
       "      <td>5.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>家和万事惊</td>\n",
       "      <td>635.2</td>\n",
       "      <td>1.5</td>\n",
       "      <td>25.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>钢铁飞龙之奥特曼崛起</td>\n",
       "      <td>595.5</td>\n",
       "      <td>1.4</td>\n",
       "      <td>3.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>海王</td>\n",
       "      <td>500.3</td>\n",
       "      <td>1.2</td>\n",
       "      <td>5.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>一条狗的回家路</td>\n",
       "      <td>360.0</td>\n",
       "      <td>0.9</td>\n",
       "      <td>4.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>掠食城市</td>\n",
       "      <td>356.6</td>\n",
       "      <td>0.9</td>\n",
       "      <td>3.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div><br>二零一九年第四周大黄蜂和密室逃生这两部影片的票房总占比是多少呀<br>sel: [2]<br>agg: ['SUM']<br>cond_conn_op: 'or'<br>conds: [[0, '==', '大黄蜂'], [0, '==', '密室逃生']]"
      ],
      "text/plain": [
       "<nl2sql.utils.Query at 0x7f40f3efc668>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(41522, 4396, 4086)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_data), len(val_data), len(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization and Label Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_brackets(s):\n",
    "    '''\n",
    "    Remove brackets [] () from text\n",
    "    '''\n",
    "    return re.sub(r'[\\(\\（].*[\\)\\）]', '', s)\n",
    "\n",
    "class QueryTokenizer(MultiSentenceTokenizer):\n",
    "    \"\"\"\n",
    "    Tokenize query (question + table header) and encode to integer sequence.\n",
    "    Using reserved tokens [unused11] and [unused12] for classification\n",
    "    \"\"\"\n",
    "    \n",
    "    col_type_token_dict = {'text': '[unused11]', 'real': '[unused12]'}\n",
    "    \n",
    "    def tokenize(self, query: Query, col_orders=None):\n",
    "        \"\"\"\n",
    "        Tokenize quesiton and columns and concatenate.\n",
    "        \n",
    "        Parameters:\n",
    "        query (Query): A query object contains question and table\n",
    "        col_orders (list or numpy.array): For re-ordering the header columns\n",
    "        \n",
    "        Returns:\n",
    "        token_idss: token ids for bert encoder\n",
    "        segment_ids: segment ids for bert encoder\n",
    "        header_ids: positions of columns\n",
    "        \"\"\"\n",
    "        \n",
    "        question_tokens = [self._token_cls] + self._tokenize(query.question.text)\n",
    "        header_tokens = []\n",
    "        \n",
    "        if col_orders is None:\n",
    "            col_orders = np.arange(len(query.table.header))\n",
    "        \n",
    "        header = [query.table.header[i] for i in col_orders]\n",
    "        \n",
    "        for col_name, col_type in header:\n",
    "            col_type_token = self.col_type_token_dict[col_type]\n",
    "            col_name = remove_brackets(col_name)\n",
    "            col_name_tokens = self._tokenize(col_name)\n",
    "            col_tokens = [col_type_token] + col_name_tokens\n",
    "            header_tokens.append(col_tokens)\n",
    "            \n",
    "        all_tokens = [question_tokens] + header_tokens\n",
    "        return self._pack(*all_tokens)\n",
    "    \n",
    "    def encode(self, query:Query, col_orders=None):\n",
    "        tokens, tokens_lens = self.tokenize(query, col_orders)\n",
    "        token_ids = self._convert_tokens_to_ids(tokens)\n",
    "        segment_ids = [0] * len(token_ids)\n",
    "        header_indices = np.cumsum(tokens_lens)\n",
    "        return token_ids, segment_ids, header_indices[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_dict = load_vocabulary(paths.vocab)\n",
    "query_tokenizer = QueryTokenizer(token_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QueryTokenizer\n",
      "\n",
      "Input Question:\n",
      "二零一九年第四周大黄蜂和密室逃生这两部影片的票房总占比是多少呀\n",
      "\n",
      "Input Header:\n",
      "影片名称(text) | 周票房（万）(real) | 票房占比（%）(real) | 场均人次(real)\n",
      "\n",
      "Output Tokens:\n",
      "[CLS] 二 零 一 九 年 第 四 周 大 黄 蜂 和 密 室 逃 生 这 两 部 影 片 的 票 房 总 占 比 是 多 少 呀 [SEP] [unused11] 影 片 名 称 [SEP] [unused12] 周 票 房 [SEP] [unused12] 票 房 占 比 [SEP] [unused12] 场 均 人 次 [SEP]\n",
      "\n",
      "Output token_ids:\n",
      "[101, 753, 7439, 671, 736, 2399, 5018, 1724, 1453, 1920, 7942, 6044, 1469, 2166, 2147, 6845, 4495, 6821, 697, 6956, 2512, 4275, 4638, 4873, 2791, 2600, 1304, 3683, 3221, 1914, 2208, 1435, 102, 11, 2512, 4275, 1399, 4917, 102, 12, 1453, 4873, 2791, 102, 12, 4873, 2791, 1304, 3683, 102, 12, 1767, 1772, 782, 3613, 102]\n",
      "Output segment_ids:\n",
      "[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "Output header_ids:\n",
      "[33 39 44 50]\n"
     ]
    }
   ],
   "source": [
    "print('QueryTokenizer\\n')\n",
    "print('Input Question:\\n{}\\n'.format(sample_query.question))\n",
    "print('Input Header:\\n{}\\n'.format(sample_query.table.header))\n",
    "print('Output Tokens:\\n{}\\n'.format(' '.join(query_tokenizer.tokenize(sample_query)[0])))\n",
    "print('Output token_ids:\\n{}\\nOutput segment_ids:\\n{}\\nOutput header_ids:\\n{}'\n",
    "      .format(*query_tokenizer.encode(sample_query)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SqlLabelEncoder:\n",
    "    \"\"\"\n",
    "    Convert SQL object into training labels.\n",
    "    \"\"\"\n",
    "    def encode(self, sql: SQL, num_cols):\n",
    "        cond_conn_op_label = sql.cond_conn_op\n",
    "        \n",
    "        sel_agg_label = np.ones(num_cols, dtype='int32') * len(SQL.agg_sql_dict)\n",
    "        for col_id, agg_op in zip(sql.sel, sql.agg):\n",
    "            if col_id < num_cols:\n",
    "                sel_agg_label[col_id] = agg_op\n",
    "            \n",
    "        cond_op_label = np.ones(num_cols, dtype='int32') * len(SQL.op_sql_dict)\n",
    "        for col_id, cond_op, _ in sql.conds:\n",
    "            if col_id < num_cols:\n",
    "                cond_op_label[col_id] = cond_op\n",
    "            \n",
    "        return cond_conn_op_label, sel_agg_label, cond_op_label\n",
    "    \n",
    "    def decode(self, cond_conn_op_label, sel_agg_label, cond_op_label):\n",
    "        cond_conn_op = int(cond_conn_op_label)\n",
    "        sel, agg, conds = [], [], []\n",
    "\n",
    "        for col_id, (agg_op, cond_op) in enumerate(zip(sel_agg_label, cond_op_label)):\n",
    "            if agg_op < len(SQL.agg_sql_dict):\n",
    "                sel.append(col_id)\n",
    "                agg.append(int(agg_op))\n",
    "            if cond_op < len(SQL.op_sql_dict):\n",
    "                conds.append([col_id, int(cond_op)])\n",
    "        return {\n",
    "            'sel': sel,\n",
    "            'agg': agg,\n",
    "            'cond_conn_op': cond_conn_op,\n",
    "            'conds': conds\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_encoder = SqlLabelEncoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cond_conn_op': 2,\n",
       " 'sel': [2],\n",
       " 'agg': [5],\n",
       " 'conds': [[0, 2, '大黄蜂'], [0, 2, '密室逃生']]}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dict(sample_query.sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, array([6, 6, 5, 6], dtype=int32), array([2, 4, 4, 4], dtype=int32))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_encoder.encode(sample_query.sql, num_cols=len(sample_query.table.header))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sel': [2], 'agg': [5], 'cond_conn_op': 2, 'conds': [[0, 2]]}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_encoder.decode(*label_encoder.encode(sample_query.sql, num_cols=len(sample_query.table.header)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build DataSequence for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataSequence(Sequence):\n",
    "    \"\"\"\n",
    "    Generate training data in batches\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 data, \n",
    "                 tokenizer, \n",
    "                 label_encoder, \n",
    "                 is_train=True, \n",
    "                 max_len=160, \n",
    "                 batch_size=32, \n",
    "                 shuffle=True, \n",
    "                 shuffle_header=True, \n",
    "                 global_indices=None):\n",
    "        \n",
    "        self.data = data\n",
    "        self.batch_size = batch_size\n",
    "        self.tokenizer = tokenizer\n",
    "        self.label_encoder = label_encoder\n",
    "        self.shuffle = shuffle\n",
    "        self.shuffle_header = shuffle_header\n",
    "        self.is_train = is_train\n",
    "        self.max_len = max_len\n",
    "        \n",
    "        if global_indices is None:\n",
    "            self._global_indices = np.arange(len(data))\n",
    "        else:\n",
    "            self._global_indices = global_indices\n",
    "\n",
    "        if shuffle:\n",
    "            np.random.shuffle(self._global_indices)\n",
    "    \n",
    "    def _pad_sequences(self, seqs, max_len=None):\n",
    "        padded = pad_sequences(seqs, maxlen=None, padding='post', truncating='post')\n",
    "        if max_len is not None:\n",
    "            padded = padded[:, :max_len]\n",
    "        return padded\n",
    "    \n",
    "    def __getitem__(self, batch_id):\n",
    "        batch_data_indices = \\\n",
    "            self._global_indices[batch_id * self.batch_size: (batch_id + 1) * self.batch_size]\n",
    "        batch_data = [self.data[i] for i in batch_data_indices]\n",
    "        \n",
    "        TOKEN_IDS, SEGMENT_IDS = [], []\n",
    "        HEADER_IDS, HEADER_MASK = [], []\n",
    "        \n",
    "        COND_CONN_OP = []\n",
    "        SEL_AGG = []\n",
    "        COND_OP = []\n",
    "        \n",
    "        for query in batch_data:\n",
    "            question = query.question.text\n",
    "            table = query.table\n",
    "            \n",
    "            col_orders = np.arange(len(table.header))\n",
    "            if self.shuffle_header:\n",
    "                np.random.shuffle(col_orders)\n",
    "            \n",
    "            token_ids, segment_ids, header_ids = self.tokenizer.encode(query, col_orders)\n",
    "            header_ids = [hid for hid in header_ids if hid < self.max_len]\n",
    "            header_mask = [1] * len(header_ids)\n",
    "            col_orders = col_orders[: len(header_ids)]\n",
    "            \n",
    "            TOKEN_IDS.append(token_ids)\n",
    "            SEGMENT_IDS.append(segment_ids)\n",
    "            HEADER_IDS.append(header_ids)\n",
    "            HEADER_MASK.append(header_mask)\n",
    "            \n",
    "            if not self.is_train:\n",
    "                continue\n",
    "            sql = query.sql\n",
    "            \n",
    "            cond_conn_op, sel_agg, cond_op = self.label_encoder.encode(sql, num_cols=len(table.header))\n",
    "            \n",
    "            sel_agg = sel_agg[col_orders]\n",
    "            cond_op = cond_op[col_orders]\n",
    "            \n",
    "            COND_CONN_OP.append(cond_conn_op)\n",
    "            SEL_AGG.append(sel_agg)\n",
    "            COND_OP.append(cond_op)\n",
    "            \n",
    "        TOKEN_IDS = self._pad_sequences(TOKEN_IDS, max_len=self.max_len)\n",
    "        SEGMENT_IDS = self._pad_sequences(SEGMENT_IDS, max_len=self.max_len)\n",
    "        HEADER_IDS = self._pad_sequences(HEADER_IDS)\n",
    "        HEADER_MASK = self._pad_sequences(HEADER_MASK)\n",
    "        \n",
    "        inputs = {\n",
    "            'input_token_ids': TOKEN_IDS,\n",
    "            'input_segment_ids': SEGMENT_IDS,\n",
    "            'input_header_ids': HEADER_IDS,\n",
    "            'input_header_mask': HEADER_MASK\n",
    "        }\n",
    "        \n",
    "        if self.is_train:\n",
    "            SEL_AGG = self._pad_sequences(SEL_AGG)\n",
    "            SEL_AGG = np.expand_dims(SEL_AGG, axis=-1)\n",
    "            COND_CONN_OP = np.expand_dims(COND_CONN_OP, axis=-1)\n",
    "            COND_OP = self._pad_sequences(COND_OP)\n",
    "            COND_OP = np.expand_dims(COND_OP, axis=-1)\n",
    "\n",
    "            outputs = {\n",
    "                'output_sel_agg': SEL_AGG,\n",
    "                'output_cond_conn_op': COND_CONN_OP,\n",
    "                'output_cond_op': COND_OP\n",
    "            }\n",
    "            return inputs, outputs\n",
    "        else:\n",
    "            return inputs\n",
    "    \n",
    "    def __len__(self):\n",
    "        return math.ceil(len(self.data) / self.batch_size)\n",
    "    \n",
    "    def on_epoch_end(self):\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(self._global_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_seq = DataSequence(train_data, query_tokenizer, label_encoder, shuffle=False, max_len=160, batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_token_ids : shape(2, 57)\n",
      "[[ 101  753 7439  671  736 2399 5018 1724 1453 1920 7942 6044 1469 2166\n",
      "  2147 6845 4495 6821  697 6956 2512 4275 4638 4873 2791 2600 1304 3683\n",
      "  3221 1914 2208 1435  102   12 1767 1772  782 3613  102   12 1453 4873\n",
      "  2791  102   12 4873 2791 1304 3683  102   11 2512 4275 1399 4917  102\n",
      "     0]\n",
      " [ 101  872 1962 8024  872 4761 6887  791 2399 5018 1724 1453 2166 2147\n",
      "  6845 4495 8024 6820 3300 6929 6956 1920 7942 6044 2124  812 4873 2791\n",
      "  2600 4638 1304 3683 1408  102   12 4873 2791 1304 3683  102   11 2512\n",
      "  4275 1399 4917  102   12 1453 4873 2791  102   12 1767 1772  782 3613\n",
      "   102]]\n",
      "input_segment_ids : shape(2, 57)\n",
      "[[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]\n",
      "input_header_ids : shape(2, 4)\n",
      "[[33 39 44 50]\n",
      " [34 40 46 51]]\n",
      "input_header_mask : shape(2, 4)\n",
      "[[1 1 1 1]\n",
      " [1 1 1 1]]\n",
      "output_sel_agg : shape(2, 4, 1)\n",
      "[[[6]\n",
      "  [6]\n",
      "  [5]\n",
      "  [6]]\n",
      "\n",
      " [[5]\n",
      "  [6]\n",
      "  [6]\n",
      "  [6]]]\n",
      "output_cond_conn_op : shape(2, 1)\n",
      "[[2]\n",
      " [2]]\n",
      "output_cond_op : shape(2, 4, 1)\n",
      "[[[4]\n",
      "  [4]\n",
      "  [4]\n",
      "  [2]]\n",
      "\n",
      " [[4]\n",
      "  [2]\n",
      "  [4]\n",
      "  [4]]]\n"
     ]
    }
   ],
   "source": [
    "sample_batch_inputs, sample_batch_outputs = train_seq[0]\n",
    "for name, data in sample_batch_inputs.items():\n",
    "    print('{} : shape{}'.format(name, data.shape))\n",
    "    print(data)\n",
    "    \n",
    "for name, data in sample_batch_outputs.items():\n",
    "    print('{} : shape{}'.format(name, data.shape))\n",
    "    print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# output sizes\n",
    "num_sel_agg = len(SQL.agg_sql_dict) + 1\n",
    "num_cond_op = len(SQL.op_sql_dict) + 1\n",
    "num_cond_conn_op = len(SQL.conn_sql_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seq_gather(x):\n",
    "    seq, idxs = x\n",
    "    idxs = K.cast(idxs, 'int32')\n",
    "    return K.tf.batch_gather(seq, idxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0925 07:42:10.219552 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:517: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.\n",
      "\n",
      "W0925 07:42:10.240487 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:4140: The name tf.random_uniform is deprecated. Please use tf.random.uniform instead.\n",
      "\n",
      "W0925 07:42:10.288790 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:131: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
      "\n",
      "W0925 07:42:10.289687 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:133: The name tf.placeholder_with_default is deprecated. Please use tf.compat.v1.placeholder_with_default instead.\n",
      "\n",
      "W0925 07:42:10.297437 139919471556416 deprecation.py:506] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:3447: calling dropout (from tensorflow.python.ops.nn_ops) with keep_prob is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `rate` instead of `keep_prob`. Rate should be set to `rate = 1 - keep_prob`.\n",
      "W0925 07:42:10.325663 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/backend/tensorflow_backend.py:4187: The name tf.truncated_normal is deprecated. Please use tf.random.truncated_normal instead.\n",
      "\n",
      "W0925 07:42:34.697302 139919471556416 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/dispatch.py:180: batch_gather (from tensorflow.python.ops.array_ops) is deprecated and will be removed after 2017-10-25.\n",
      "Instructions for updating:\n",
      "`tf.batch_gather` is deprecated, please use `tf.gather` with `batch_dims=-1` instead.\n"
     ]
    }
   ],
   "source": [
    "bert_model = load_trained_model_from_checkpoint(paths.config, paths.checkpoint, seq_len=None)\n",
    "for l in bert_model.layers:\n",
    "    l.trainable = True\n",
    "    \n",
    "inp_token_ids = Input(shape=(None,), name='input_token_ids', dtype='int32')\n",
    "inp_segment_ids = Input(shape=(None,), name='input_segment_ids', dtype='int32')\n",
    "inp_header_ids = Input(shape=(None,), name='input_header_ids', dtype='int32')\n",
    "inp_header_mask = Input(shape=(None, ), name='input_header_mask')\n",
    "\n",
    "x = bert_model([inp_token_ids, inp_segment_ids]) # (None, seq_len, 768)\n",
    "\n",
    "# predict cond_conn_op\n",
    "x_for_cond_conn_op = Lambda(lambda x: x[:, 0])(x) # (None, 768)\n",
    "p_cond_conn_op = Dense(num_cond_conn_op, activation='softmax', name='output_cond_conn_op')(x_for_cond_conn_op)\n",
    "\n",
    "# predict sel_agg\n",
    "x_for_header = Lambda(seq_gather, name='header_seq_gather')([x, inp_header_ids]) # (None, h_len, 768)\n",
    "header_mask = Lambda(lambda x: K.expand_dims(x, axis=-1))(inp_header_mask) # (None, h_len, 1)\n",
    "\n",
    "x_for_header = Multiply()([x_for_header, header_mask])\n",
    "x_for_header = Masking()(x_for_header)\n",
    "\n",
    "p_sel_agg = Dense(num_sel_agg, activation='softmax', name='output_sel_agg')(x_for_header)\n",
    "\n",
    "x_for_cond_op = Concatenate(axis=-1)([x_for_header, p_sel_agg])\n",
    "p_cond_op = Dense(num_cond_op, activation='softmax', name='output_cond_op')(x_for_cond_op)\n",
    "\n",
    "\n",
    "model = Model(\n",
    "    [inp_token_ids, inp_segment_ids, inp_header_ids, inp_header_mask],\n",
    "    [p_cond_conn_op, p_sel_agg, p_cond_op]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using 2 gpus\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0925 07:42:44.093672 139919471556416 deprecation_wrapper.py:118] From /usr/local/lib/python3.6/dist-packages/keras/optimizers.py:790: The name tf.train.Optimizer is deprecated. Please use tf.compat.v1.train.Optimizer instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "NUM_GPUS = 2\n",
    "if NUM_GPUS > 1:\n",
    "    print('using {} gpus'.format(NUM_GPUS))\n",
    "    model = multi_gpu_model(model, gpus=NUM_GPUS)\n",
    "\n",
    "learning_rate = 1e-5\n",
    "\n",
    "model.compile(\n",
    "    loss='sparse_categorical_crossentropy',\n",
    "    optimizer=RAdam(lr=learning_rate)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def outputs_to_sqls(preds_cond_conn_op, preds_sel_agg, preds_cond_op, header_lens, label_encoder):\n",
    "    \"\"\"\n",
    "    Generate sqls from model outputs\n",
    "    \"\"\"\n",
    "    preds_cond_conn_op = np.argmax(preds_cond_conn_op, axis=-1)\n",
    "    preds_cond_op = np.argmax(preds_cond_op, axis=-1)\n",
    "\n",
    "    sqls = []\n",
    "    \n",
    "    for cond_conn_op, sel_agg, cond_op, header_len in zip(preds_cond_conn_op, \n",
    "                                                          preds_sel_agg, \n",
    "                                                          preds_cond_op, \n",
    "                                                          header_lens):\n",
    "        sel_agg = sel_agg[:header_len]\n",
    "        # force to select at least one column for agg\n",
    "        sel_agg[sel_agg == sel_agg[:, :-1].max()] = 1\n",
    "        sel_agg = np.argmax(sel_agg, axis=-1)\n",
    "        \n",
    "        sql = label_encoder.decode(cond_conn_op, sel_agg, cond_op)\n",
    "        sql['conds'] = [cond for cond in sql['conds'] if cond[0] < header_len]\n",
    "        \n",
    "        sel = []\n",
    "        agg = []\n",
    "        for col_id, agg_op in zip(sql['sel'], sql['agg']):\n",
    "            if col_id < header_len:\n",
    "                sel.append(col_id)\n",
    "                agg.append(agg_op)\n",
    "                \n",
    "        sql['sel'] = sel\n",
    "        sql['agg'] = agg\n",
    "        sqls.append(sql)\n",
    "    return sqls\n",
    "\n",
    "class EvaluateCallback(Callback):\n",
    "    def __init__(self, val_dataseq):\n",
    "        self.val_dataseq = val_dataseq\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        pred_sqls = []\n",
    "        for batch_data in self.val_dataseq:\n",
    "            header_lens = np.sum(batch_data['input_header_mask'], axis=-1)\n",
    "            preds_cond_conn_op, preds_sel_agg, preds_cond_op = self.model.predict_on_batch(batch_data)\n",
    "            sqls = outputs_to_sqls(preds_cond_conn_op, preds_sel_agg, preds_cond_op, \n",
    "                                   header_lens, val_dataseq.label_encoder)\n",
    "            pred_sqls += sqls\n",
    "            \n",
    "        conn_correct = 0\n",
    "        agg_correct = 0\n",
    "        conds_correct = 0\n",
    "        conds_col_id_correct = 0\n",
    "        all_correct = 0\n",
    "        num_queries = len(self.val_dataseq.data)\n",
    "        \n",
    "        true_sqls = [query.sql for query in self.val_dataseq.data]\n",
    "        for pred_sql, true_sql in zip(pred_sqls, true_sqls):\n",
    "            n_correct = 0\n",
    "            if pred_sql['cond_conn_op'] == true_sql.cond_conn_op:\n",
    "                conn_correct += 1\n",
    "                n_correct += 1\n",
    "            \n",
    "            pred_aggs = set(zip(pred_sql['sel'], pred_sql['agg']))\n",
    "            true_aggs = set(zip(true_sql.sel, true_sql.agg))\n",
    "            if pred_aggs == true_aggs:\n",
    "                agg_correct += 1\n",
    "                n_correct += 1\n",
    "\n",
    "            pred_conds = set([(cond[0], cond[1]) for cond in pred_sql['conds']])\n",
    "            true_conds = set([(cond[0], cond[1]) for cond in true_sql.conds])\n",
    "\n",
    "            if pred_conds == true_conds:\n",
    "                conds_correct += 1\n",
    "                n_correct += 1\n",
    "   \n",
    "            pred_conds_col_ids = set([cond[0] for cond in pred_sql['conds']])\n",
    "            true_conds_col_ids = set([cond[0] for cond in true_sql['conds']])\n",
    "            if pred_conds_col_ids == true_conds_col_ids:\n",
    "                conds_col_id_correct += 1\n",
    "            \n",
    "            if n_correct == 3:\n",
    "                all_correct += 1\n",
    "\n",
    "        print('conn_acc: {}'.format(conn_correct / num_queries))\n",
    "        print('agg_acc: {}'.format(agg_correct / num_queries))\n",
    "        print('conds_acc: {}'.format(conds_correct / num_queries))\n",
    "        print('conds_col_id_acc: {}'.format(conds_col_id_correct / num_queries))\n",
    "        print('total_acc: {}'.format(all_correct / num_queries))\n",
    "        \n",
    "        logs['val_tot_acc'] = all_correct / num_queries\n",
    "        logs['conn_acc'] = conn_correct / num_queries\n",
    "        logs['conds_acc'] = conds_correct / num_queries\n",
    "        logs['conds_col_id_acc'] = conds_col_id_correct / num_queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = NUM_GPUS * 32\n",
    "num_epochs = 30\n",
    "\n",
    "train_dataseq = DataSequence(\n",
    "    data=train_data,\n",
    "    tokenizer=query_tokenizer,\n",
    "    label_encoder=label_encoder,\n",
    "    shuffle_header=False,\n",
    "    is_train=True, \n",
    "    max_len=160, \n",
    "    batch_size=batch_size\n",
    ")\n",
    "\n",
    "val_dataseq = DataSequence(\n",
    "    data=val_data, \n",
    "    tokenizer=query_tokenizer,\n",
    "    label_encoder=label_encoder,\n",
    "    is_train=False, \n",
    "    shuffle_header=False,\n",
    "    max_len=160, \n",
    "    shuffle=False,\n",
    "    batch_size=batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'task1_best_model.h5'\n",
    "callbacks = [\n",
    "    EvaluateCallback(val_dataseq),\n",
    "    ModelCheckpoint(filepath=model_path, \n",
    "                    monitor='val_tot_acc', \n",
    "                    mode='max', \n",
    "                    save_best_only=True, \n",
    "                    save_weights_only=True)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0925 07:43:15.483318 139919471556416 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/math_grad.py:1251: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "649/649 [==============================] - 434s 669ms/step - loss: 1.6993 - output_cond_conn_op_loss: 0.6747 - output_sel_agg_loss: 0.4985 - output_cond_op_loss: 0.5261\n",
      "conn_acc: 0.9103730664240218\n",
      "agg_acc: 0.591901728844404\n",
      "conds_acc: 0.3964968152866242\n",
      "conds_col_id_acc: 0.4294813466787989\n",
      "total_acc: 0.2456778889899909\n",
      "Epoch 2/3\n",
      "649/649 [==============================] - 429s 661ms/step - loss: 0.5754 - output_cond_conn_op_loss: 0.1915 - output_sel_agg_loss: 0.1659 - output_cond_op_loss: 0.2180\n",
      "conn_acc: 0.9520018198362148\n",
      "agg_acc: 0.8328025477707006\n",
      "conds_acc: 0.6808462238398544\n",
      "conds_col_id_acc: 0.6951774340309372\n",
      "total_acc: 0.5643767060964513\n",
      "Epoch 3/3\n",
      "649/649 [==============================] - 430s 662ms/step - loss: 0.3120 - output_cond_conn_op_loss: 0.1096 - output_sel_agg_loss: 0.0793 - output_cond_op_loss: 0.1231\n",
      "conn_acc: 0.9565514103730665\n",
      "agg_acc: 0.8971792538671519\n",
      "conds_acc: 0.7686533212010919\n",
      "conds_col_id_acc: 0.7807097361237488\n",
      "total_acc: 0.6822111010009099\n"
     ]
    }
   ],
   "source": [
    "history = model.fit_generator(train_dataseq, epochs=num_epochs, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make prediction for task1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataseq = DataSequence(\n",
    "    data=test_data, \n",
    "    tokenizer=query_tokenizer,\n",
    "    label_encoder=label_encoder,\n",
    "    is_train=False, \n",
    "    shuffle_header=False,\n",
    "    max_len=160, \n",
    "    shuffle=False,\n",
    "    batch_size=batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "901d16d8661d4dcd9e7d23488a1bd4b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=64), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "pred_sqls = []\n",
    "\n",
    "for batch_data in tqdm(test_dataseq):\n",
    "    header_lens = np.sum(batch_data['input_header_mask'], axis=-1)\n",
    "    preds_cond_conn_op, preds_sel_agg, preds_cond_op = model.predict_on_batch(batch_data)\n",
    "    sqls = outputs_to_sqls(preds_cond_conn_op, preds_sel_agg, preds_cond_op, \n",
    "                           header_lens, val_dataseq.label_encoder)\n",
    "    pred_sqls += sqls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_output_file = 'task1_output.json'\n",
    "with open(task1_output_file, 'w') as f:\n",
    "    for sql in pred_sqls:\n",
    "        json_str = json.dumps(sql, ensure_ascii=False)\n",
    "        f.write(json_str + '\\n')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
